#!/usr/bin/env python3
import subprocess
import os

def get_current_crontab():
    try:
        result = subprocess.run(['crontab', '-l'],
                              stdout=subprocess.PIPE,
                              stderr=subprocess.PIPE,
                              universal_newlines=True)
        return result.stdout
    except:
        return ""

def add_auto_check_cron():
    cron_entry = "0 * * * * cd /root/chaosblade_manager && python3 server_stress_advanced.py status 2>&1 | grep -q '活跃实验数: 0' && python3 server_stress_advanced.py cpu --cpu-percent 30"

    current_cron = get_current_crontab()

    if 'chaosblade_manager' in current_cron and 'server_stress_advanced.py' in current_cron:
        print("定时任务已存在！")
        print("\n当前定时任务：")
        print(current_cron)
        return False

    new_cron = current_cron.strip() + '\n' + cron_entry + '\n'

    try:
        process = subprocess.Popen(['crontab', '-'], stdin=subprocess.PIPE, universal_newlines=True)
        process.communicate(input=new_cron)

        if process.returncode == 0:
            print("✓ 定时任务添加成功！")
            print("\n定时规则：每小时检查一次，如果压测停止则自动启动（CPU 30%）")
            print("作用：防止服务器重启后压测失效")
            print("\n查看定时任务：crontab -l")
            return True
        else:
            print("✗ 定时任务添加失败")
            return False
    except Exception as e:
        print(f"✗ 错误: {e}")
        return False

def remove_auto_check_cron():
    current_cron = get_current_crontab()

    if 'chaosblade_manager' not in current_cron:
        print("未找到相关定时任务")
        return False

    lines = current_cron.split('\n')
    new_lines = [line for line in lines if 'chaosblade_manager' not in line or 'server_stress_advanced.py' not in line]
    new_cron = '\n'.join(new_lines)

    try:
        process = subprocess.Popen(['crontab', '-'], stdin=subprocess.PIPE, universal_newlines=True)
        process.communicate(input=new_cron)

        if process.returncode == 0:
            print("✓ 定时任务已移除")
            return True
        else:
            print("✗ 移除失败")
            return False
    except Exception as e:
        print(f"✗ 错误: {e}")
        return False

def show_crontab():
    cron = get_current_crontab()
    if cron:
        print("当前定时任务：")
        print("="*80)
        print(cron)
        print("="*80)
    else:
        print("暂无定时任务")

if __name__ == '__main__':
    import sys

    if len(sys.argv) < 2:
        print("ChaosBlade 定时任务管理工具")
        print("\n用法:")
        print("  python3 cron_manager.py add     - 添加自动检查定时任务")
        print("  python3 cron_manager.py remove  - 移除定时任务")
        print("  python3 cron_manager.py show    - 查看定时任务")
        print("\n推荐：添加定时任务，防止服务器重启后压测失效")
        sys.exit(1)

    action = sys.argv[1]

    if action == 'add':
        add_auto_check_cron()
    elif action == 'remove':
        remove_auto_check_cron()
    elif action == 'show':
        show_crontab()
    else:
        print("未知操作，使用: add/remove/show")
